{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Knowledge Distillation for Computer Vision"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "知識の蒸留は、より大規模で複雑なモデル (教師) からより小規模で単純なモデル (生徒) に知識を伝達するために使用される手法です。あるモデルから別のモデルに知識を抽出するには、特定のタスク (この場合は画像分類) でトレーニングされた事前トレーニング済み教師モデルを取得し、画像分類でトレーニングされる生徒モデルをランダムに初期化します。次に、学生モデルをトレーニングして、その出力と教師の出力の差を最小限に抑え、動作を模倣します。これは [Distilling the Knowledge in a Neural Network by Hinton et al](https://huggingface.co/papers/1503.02531) で最初に導入されました。このガイドでは、タスク固有の知識の蒸留を行います。これには [Beans データセット](https://huggingface.co/datasets/beans) を使用します。\n",
    "\n",
    "このガイドでは、[微調整された ViT モデル](https://huggingface.co/merve/vit-mobilenet-beans-224) (教師モデル) を抽出して [MobileNet](https://huggingface.co/google/mobilenet_v2_1.4_224) (学生モデル) 🤗 Transformers の [Trainer API](https://huggingface.co/docs/transformers/en/main_classes/trainer#trainer) を使用します。\n",
    "\n",
    "蒸留とプロセスの評価に必要なライブラリをインストールしましょう。\n",
    "\n",
    "```bash\n",
    "pip install transformers datasets accelerate tensorboard evaluate --upgrade\n",
    "```\n",
    "\n",
    "この例では、教師モデルとして`merve/beans-vit-224`モデルを使用しています。これは、Bean データセットに基づいて微調整された`google/vit-base-patch16-224-in21k`に基づく画像分類モデルです。このモデルをランダムに初期化された MobileNetV2 に抽出します。\n",
    "\n",
    "次に、データセットをロードします。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "dataset = load_dataset(\"beans\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "この場合、同じ解像度で同じ出力が返されるため、どちらのモデルの画像プロセッサも使用できます。 `dataset`の`map()`メソッドを使用して、データセットのすべての分割に前処理を適用します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoImageProcessor\n",
    "teacher_processor = AutoImageProcessor.from_pretrained(\"merve/beans-vit-224\")\n",
    "\n",
    "def process(examples):\n",
    "    processed_inputs = teacher_processor(examples[\"image\"])\n",
    "    return processed_inputs\n",
    "\n",
    "processed_datasets = dataset.map(process, batched=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "基本的に、我々は生徒モデル（ランダムに初期化されたMobileNet）が教師モデル（微調整されたビジョン変換器）を模倣することを望む。これを実現するために、まず教師と生徒からロジット出力を得る。次に、それぞれのソフトターゲットの重要度を制御するパラメータ`temperature`で分割する。`lambda`と呼ばれるパラメータは蒸留ロスの重要度を量る。この例では、`temperature=5`、`lambda=0.5`とする。生徒と教師の間の発散を計算するために、Kullback-Leibler発散損失を使用します。2つのデータPとQが与えられたとき、KLダイバージェンスはQを使ってPを表現するためにどれだけの余分な情報が必要かを説明します。もし2つが同じであれば、QからPを説明するために必要な他の情報はないので、それらのKLダイバージェンスはゼロになります。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import TrainingArguments, Trainer\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "\n",
    "class ImageDistilTrainer(Trainer):\n",
    "    def __init__(self, *args, teacher_model=None, **kwargs):\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.teacher = teacher_model\n",
    "        self.student = student_model\n",
    "        self.loss_function = nn.KLDivLoss(reduction=\"batchmean\")\n",
    "        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "        self.teacher.to(device)\n",
    "        self.teacher.eval()\n",
    "        self.temperature = temperature\n",
    "        self.lambda_param = lambda_param\n",
    "\n",
    "    def compute_loss(self, student, inputs, return_outputs=False):\n",
    "        student_output = self.student(**inputs)\n",
    "\n",
    "        with torch.no_grad():\n",
    "          teacher_output = self.teacher(**inputs)\n",
    "\n",
    "        # Compute soft targets for teacher and student\n",
    "        soft_teacher = F.softmax(teacher_output.logits / self.temperature, dim=-1)\n",
    "        soft_student = F.log_softmax(student_output.logits / self.temperature, dim=-1)\n",
    "\n",
    "        # Compute the loss\n",
    "        distillation_loss = self.loss_function(soft_student, soft_teacher) * (self.temperature ** 2)\n",
    "\n",
    "        # Compute the true label loss\n",
    "        student_target_loss = student_output.loss\n",
    "\n",
    "        # Calculate final loss\n",
    "        loss = (1. - self.lambda_param) * student_target_loss + self.lambda_param * distillation_loss\n",
    "        return (loss, student_output) if return_outputs else loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "次に、Hugging Face Hub にログインして、`trainer`を通じてモデルを Hugging Face Hub にプッシュできるようにします。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import notebook_login\n",
    "\n",
    "notebook_login()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "教師モデルと生徒モデルである`TrainingArguments`を設定しましょう。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForImageClassification, MobileNetV2Config, MobileNetV2ForImageClassification\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=\"my-awesome-model\",\n",
    "    num_train_epochs=30,\n",
    "    fp16=True,\n",
    "    logging_strategy=\"epoch\",\n",
    "    eval_strategy=\"epoch\",\n",
    "    save_strategy=\"epoch\",\n",
    "    load_best_model_at_end=True,\n",
    "    metric_for_best_model=\"accuracy\",\n",
    "    report_to=\"tensorboard\",\n",
    "    push_to_hub=True,\n",
    "    hub_strategy=\"every_save\",\n",
    "    hub_model_id=repo_name,\n",
    "    )\n",
    "\n",
    "num_labels = len(processed_datasets[\"train\"].features[\"labels\"].names)\n",
    "\n",
    "# initialize models\n",
    "teacher_model = AutoModelForImageClassification.from_pretrained(\n",
    "    \"merve/beans-vit-224\",\n",
    "    num_labels=num_labels,\n",
    "    ignore_mismatched_sizes=True\n",
    ")\n",
    "\n",
    "# training MobileNetV2 from scratch\n",
    "student_config = MobileNetV2Config()\n",
    "student_config.num_labels = num_labels\n",
    "student_model = MobileNetV2ForImageClassification(student_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`compute_metrics` 関数を使用して、テスト セットでモデルを評価できます。この関数は、トレーニング プロセス中にモデルの`accuracy`と`f1`を計算するために使用されます。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import evaluate\n",
    "import numpy as np\n",
    "\n",
    "accuracy = evaluate.load(\"accuracy\")\n",
    "\n",
    "def compute_metrics(eval_pred):\n",
    "    predictions, labels = eval_pred\n",
    "    acc = accuracy.compute(references=labels, predictions=np.argmax(predictions, axis=1))\n",
    "    return {\"accuracy\": acc[\"accuracy\"]}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "定義したトレーニング引数を使用して`Trainer`を初期化しましょう。データ照合装置も初期化します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import DefaultDataCollator\n",
    "\n",
    "data_collator = DefaultDataCollator()\n",
    "trainer = ImageDistilTrainer(\n",
    "    student_model=student_model,\n",
    "    teacher_model=teacher_model,\n",
    "    training_args=training_args,\n",
    "    train_dataset=processed_datasets[\"train\"],\n",
    "    eval_dataset=processed_datasets[\"validation\"],\n",
    "    data_collator=data_collator,\n",
    "    processing_class=teacher_extractor,\n",
    "    compute_metrics=compute_metrics,\n",
    "    temperature=5,\n",
    "    lambda_param=0.5\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "これでモデルをトレーニングできるようになりました。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "テスト セットでモデルを評価できます。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.evaluate(processed_datasets[\"test\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "テスト セットでは、モデルの精度は 72% に達します。蒸留効率の健全性チェックを行うために、同じハイパーパラメータを使用して Bean データセットで MobileNet を最初からトレーニングし、テスト セットで 63% の精度を観察しました。読者の皆様には、さまざまな事前トレーニング済み教師モデル、学生アーキテクチャ、蒸留パラメータを試していただき、その結果を報告していただくようお勧めします。抽出されたモデルのトレーニング ログとチェックポイントは [このリポジトリ](https://huggingface.co/merve/vit-mobilenet-beans-224) にあり、最初からトレーニングされた MobileNetV2 はこの [リポジトリ](https://huggingface.co/merve/resnet-mobilenet-beans-5)。"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 4
}
